{-# LANGUAGE GeneralizedNewtypeDeriving,
             NamedFieldPuns,
             FlexibleInstances,
             MultiParamTypeClasses,
             DeriveDataTypeable,
             OverloadedStrings
             #-}

module Main where

import Control.Applicative
import Control.Concurrent ( forkIO, killThread )
import Control.Concurrent.STM
import Control.Monad ( filterM, void )
import Data.Attoparsec.Text
import Data.Enumerator ( Enumeratee, ($=), ($$) )
import qualified Data.Enumerator as E ( run_ )
import Data.Enumerator.Text ( enumHandle, iterHandle )
import qualified Data.Enumerator.List as E
  ( map, concatMapAccum, mapM_, generateM )
import qualified Data.Foldable as F
import Data.Function ( on )
import Data.Functor
import Data.Int ( Int64 )
import Data.List ( intercalate, maximumBy )
import Data.Map ( Map )
import qualified Data.Map as Map
import Data.Set ( Set )
import qualified Data.Set as Set
import Data.Text ( Text )
import qualified Data.Text as Text ( pack, unpack )
import Data.Text.IO ( hGetLine, hPutStr, hPutStrLn )
import Network.Server
import Network.Socket ( PortNumber )
import qualified Network.Socket as Sock
import System.Console.CmdLib
import qualified System.Console.CmdLib as Cmd
import System.IO
  ( Handle, BufferMode(LineBuffering)
  , hClose, hSetBuffering, hSetNewlineMode, universalNewlineMode
  )
import qualified System.Posix.Daemonize as Posix
import qualified Util.Enumerator as E ( until )
import Util.Monad

import Debug.Trace


-- * 'roomd' parameters and configuration

-- HARD CODED PARAMETERS
-- TODO: Read them as command arguments
_DEFAULT_ROOM_CAPACITY :: NumPlayers
_DEFAULT_ROOM_CAPACITY = 5000

_DEFAULT_BUCKET_CAPACITY :: NumPlayers 
_DEFAULT_BUCKET_CAPACITY = 5
--

data Config = 
      Config { 
        port :: Int     -- ^ Port to listen on
      , daemon :: Bool  -- ^ daemonize?
      }
    deriving (Typeable, Data, Eq)

instance Attributes Config where
  attributes _ = Cmd.group "Options"
      [ port   %> Help "Port number." %+ ArgHelp "PORT" %+ short 'p' %+ Default (4001::Int)
      , daemon %> Help "Daemonize" %+ Long ["daemonize"] %+ short 'd'
      ]

-- * Debug

_DEBUG :: Monad m => String -> m ()
_DEBUG msg = trace ("(DEBUG) " ++ msg) $ return ()

-- * Communication

type EventChan = TChan Event

class ParseProto a where
  parseText :: Text -> a

class ShowProto a where
  showText :: a -> Text


-- * Player

type PlayerId = Int64

playerIdP :: Parser PlayerId
playerIdP = decimal

type Coordinate = Double

data Facing = UP | DOWN | LEFT | RIGHT
    deriving Show

facingP :: Parser Facing
facingP =    p_facing "UP" UP
         <|> p_facing "DOWN" DOWN 
         <|> p_facing "LEFT" LEFT
         <|> p_facing "RIGHT" RIGHT
  where p_facing s c = string (Text.pack s) >> return c

data Position
  = Position {
      ppX      :: !Coordinate
    , ppY      :: !Coordinate
    , ppFacing :: !Facing
    }

instance Show Position where
  show (Position x y f) = show x ++ "," ++ show y ++ "," ++ show f

positionP :: Parser Position
positionP = do
  x <- double
  void comma
  y <- double
  void comma
  f <- facingP
  return $ Position x y f
  where comma = char ','

type Nickname = Text

type Character = Text

data Skilly
  = Skilly {
      skillyNick :: !Nickname
    , skillyChar :: !Character
    }

instance Show Skilly where
  show (Skilly nick char) = Text.unpack nick ++ "," ++ Text.unpack char

skillyP :: Parser Skilly
skillyP = do
  nick <- takeTill isComma
  comma
  char <- takeTill isEnd
  return $ Skilly nick char
  where isComma = (== ',')
        isEnd = (`elem` [' ','\n'])
        comma = char ','

data Player
  = Player {
      playerId     :: PlayerId
    , playerHandle :: Handle
    , playerChan   :: EventChan
    , playerSkilly :: TVar Skilly
    , playerPos    :: TVar Position
    }

instance Eq Player where
  (==) = (==) `on` playerId

instance Ord Player where
  compare = compare `on` playerId

mkPlayer :: PlayerId -> Handle -> Skilly -> Position -> STM Player
mkPlayer id handle skilly pos = do
  pEChan <- newTChan
  pSkilly <- newTVar skilly
  pPos <- newTVar pos
  return $ Player id handle pEChan pSkilly pPos

updateSkilly :: Player -> Skilly -> STM ()
updateSkilly Player{playerSkilly} = writeTVar playerSkilly

updatePosition :: Player -> Position -> STM ()
updatePosition Player{playerPos} = writeTVar playerPos

-- * Actions and events

type Message = Text

messageP :: Parser Message
messageP = takeTill (== '\n')

data Action
  = Asms !Message
  | Askilly !Skilly
  | Apos !Position
  | Apart

isAPart :: Action -> Bool
isAPart Apart = True
isAPart _act  = False

-- isAPos :: Action -> Maybe Position
-- isAPos (Apos pos) = Just pos
-- isAPos _act       = Nothing

actionP :: Parser Action
actionP = p_sms <|> p_skilly <|> p_pos <|> p_part
  where p_sms    = Asms  <$> ( string (Text.pack "SMS") *> space *> messageP )
        p_skilly = Askilly <$> ( string (Text.pack "SKILLY") *> space *> skillyP )
        p_pos    = Apos <$> ( string (Text.pack "POS") *> space *> positionP )
        p_part   = string (Text.pack "PART") *> pure Apart

instance ParseProto Action where
  parseText text
    = case parseOnly actionP text of
          Left err  -> error err
          Right act -> act

updatePlayer :: Action -> Player -> STM ()
updatePlayer (Askilly newSkilly) Player{playerSkilly}
  = writeTVar playerSkilly newSkilly
updatePlayer (Apos newPos) Player{playerPos}
  = writeTVar playerPos newPos
updatePlayer _other        _player
  = return ()

data Event = Ejoin PlayerId Skilly Position
           | Ejoined [(PlayerId,Skilly,Position)]
           | Eaction PlayerId Action

instance Show Event where
  show (Ejoin id skilly pos)
      = "JOIN" ++ " " ++ show id ++ " " ++ show skilly ++ " " ++ show pos
  show (Ejoined users)         = "JOINED" ++ " " ++ show_users users
    where show_users = intercalate " " . map show_user
          show_user (id,skilly,pos) = show id ++ ":" ++ show skilly ++ ":" ++ show pos
  show (Eaction id (Asms msg)) = "SMS"  ++ " " ++ show id ++ " " ++ show msg
  show (Eaction id (Askilly skilly)) 
      = "SKILLY" ++ " " ++ show id ++ " " ++ show skilly
  show (Eaction id (Apos pos)) = "POS" ++ " " ++ show id ++ " " ++ show pos
  show (Eaction id Apart)      = "PART" ++ " " ++ show id

instance ShowProto Event where
  showText = Text.pack . (++ "\n") . show


-- * Room
-- A roomd is divided into buckets for scalability.

newtype NumPlayers = NumP Int
  deriving(Eq, Ord, Show, Num)

type RoomBuckets = Map BucketId Bucket

data Room
  = Room {
      roomBuckets  :: TVar RoomBuckets
    , roomPlayers  :: TVar NumPlayers
    , roomCapacity :: !NumPlayers
    }

newRoom :: STM Room
newRoom = do
  buckets <- newTVar Map.empty
  players <- newTVar 0
  return $ Room buckets players _DEFAULT_ROOM_CAPACITY

-- ** Bucket

type BucketId = Int

  -- | A broadcast channel: what you send to a bucket-chan
  -- is sent to all its players.
type BucketChan = TVar (Map PlayerId EventChan)

newBucketChan :: STM BucketChan
newBucketChan = newTVar Map.empty

type BucketPlayers = TVar (Set Player)

data Bucket
  = Bucket {
      bucketChan      :: BucketChan
    , bucketPlayers   :: BucketPlayers
    , bucketCapacity  :: !NumPlayers
    }

newBucket :: TVar RoomBuckets -> STM Bucket
newBucket roomBuckets = do
  buckets <- readTVar roomBuckets
  bChan <- newBucketChan
  bPlayers <- newTVar Set.empty
  let bucket   = Bucket bChan bPlayers _DEFAULT_BUCKET_CAPACITY
      bId      = nextId $ Map.keysSet buckets
      buckets' = Map.insert bId bucket buckets
  writeTVar roomBuckets buckets'
  _DEBUG ("New bucket = " ++ show bId)
  return bucket
  where nextId bIds | Set.null bIds = 0
                    | otherwise     = 1 + Set.findMax bIds 

nonFullBucket :: Bucket -> STM Bool
nonFullBucket Bucket{bucketPlayers,bucketCapacity} = do
  players <- readTVar bucketPlayers
  let numPlayers = fromIntegral $ Set.size players
  return $ numPlayers < bucketCapacity

  -- | An user notifies an event to a bucket
  -- Every player inside the bucket (except the messenger) is notified.
notifyBucket :: PlayerId -> Bucket -> Event -> STM ()
notifyBucket myId Bucket{bucketChan} event
  = readTVar bucketChan >>=
        mapM_ notify_event . Map.elems . Map.delete myId
  where notify_event = (`writeTChan` event)
--       notify_event tchan = do        
--         writeTChan tchan event
--         trace ("EVENT " ++ show event ++ " wrote to TChan")

findNonFullBucket :: RoomBuckets -> STM (Maybe Bucket)
findNonFullBucket roomBuckets = do
  nonFullBuckets <- filterM nonFullBucket $ Map.elems $ roomBuckets
  case nonFullBuckets of
      []        -> return Nothing
      bucket:_  -> return $ Just bucket

addPlayer :: Player -> Bucket -> STM ()
addPlayer player@Player{playerId,playerChan} Bucket{bucketChan,bucketPlayers} = do
  inc_bucket_players
  extend_bucket_chan
  _DEBUG ("[" ++ show playerId ++ "] Added to bucket")
  where inc_bucket_players = do
          players <- readTVar bucketPlayers
          writeTVar bucketPlayers (Set.insert player players)
        extend_bucket_chan = do
          chans <- readTVar bucketChan
          let chans' = Map.insert playerId playerChan chans
          _DEBUG ("[" ++ show playerId ++ "] bucket-chan = " ++ show (Map.keys chans'))
          writeTVar bucketChan chans'

delPlayer :: Player -> Bucket -> STM ()
delPlayer player@Player{playerId} Bucket{bucketChan,bucketPlayers} = do
  dec_bucket_players
  restrict_bucket_chan
  where dec_bucket_players = do
          players <- readTVar bucketPlayers
          writeTVar bucketPlayers (Set.delete player players)
        restrict_bucket_chan = do
          chans <- readTVar bucketChan  
          let chans' = Map.delete playerId chans
          writeTVar bucketChan chans'

  -- TODO: Check room capacity
assignBucket :: Player -> Room -> STM Bucket
assignBucket player Room{roomBuckets,roomPlayers} = do
  buckets <- readTVar roomBuckets
  players <- readTVar roomPlayers
  bucket <- findNonFullBucket buckets `or_else` newBucket roomBuckets
  addPlayer player bucket
  return bucket
  where or_else m_mb_a mb = do
          mb_a <- m_mb_a
          case mb_a of
              Just a  -> return a
              Nothing -> mb


-- * Server core

data Session
  = Session {
      myPlayer     :: Player
    , myBucket     :: Bucket
    , myEndMark    :: TMVar ()
    }

registerPlayer :: Room -> PlayerId -> Skilly -> Position -> Handle -> IO Session
registerPlayer room pId skilly pos pHandle = do
  player <- atomically $ mkPlayer pId pHandle skilly pos
  bucket <- atomically $ assignBucket player room
  endMark <- newEmptyTMVarIO
  return $ Session player bucket endMark

data JoinRequest = JoinReq PlayerId Skilly Position 

joinRequestP :: Parser JoinRequest
joinRequestP = do
  string (Text.pack "JOIN")
  space
  pId <- decimal
  space
  skilly <- skillyP
  space
  pos <- positionP
  return $ JoinReq pId skilly pos

instance ParseProto JoinRequest where
  parseText text
    = case parseOnly joinRequestP text of
          Left err   -> error err
          Right jreq -> jreq

readBucketEvents :: Session -> IO ()
readBucketEvents session@Session{myPlayer,myBucket,myEndMark}
  = let eventsEnum = E.generateM gen_enum $= E.map showText
      in E.run_ (eventsEnum $$ iterHandle myHandle)
  where myId        = playerId myPlayer
        myHandle    = playerHandle myPlayer
        myEventChan = playerChan myPlayer
        read_event = Just <$> readTChan myEventChan
        stop_reading = do () <- takeTMVar myEndMark
                          return Nothing
        gen_enum = atomically $ read_event `orElse` stop_reading

notifyMyActions :: Session -> IO ()
notifyMyActions session@Session{myPlayer,myBucket,myEndMark}
  = let eventsEnum = enumHandle myHandle $= E.map parseText $= E.until isAPart
     in do E.run_ (eventsEnum $$ E.mapM_ handle_action)
           notifyPart
  where myId     = playerId myPlayer
        myHandle = playerHandle myPlayer
        handle_action :: Action -> IO ()
        handle_action act = atomically $ do
            updatePlayer act myPlayer
            notifyBucket myId myBucket (Eaction myId act)
--         parseTextDebug t = trace ("[" ++ show myId ++ "] SOCKET RAW MSG: " ++ show t) $ parseText t
        notifyPart = atomically $ notifyBucket myId myBucket (Eaction myId Apart)

handlePlayer :: Room -> ServerRoutine
handlePlayer room (pHandle,_,_) = do
  _DEBUG "New connection"
    -- every command ends with a newline character
  hSetBuffering pHandle LineBuffering
  hSetNewlineMode pHandle universalNewlineMode
  JoinReq playerId skilly pos <- readJoinRequest pHandle
  session@Session{myPlayer,myBucket}
    <- registerPlayer room playerId skilly pos pHandle
  notifyJoin myPlayer myBucket
  _DEBUG ("New player " ++ show playerId)
  void $ forkIO (readBucketEvents session)
  notifyMyActions session
    -- when 'notifyMyActions' returns we stop reading events
  stopReadingEvents session
  atomically $ delPlayer myPlayer myBucket
  _DEBUG ("[" ++ show playerId ++ "] Logout")
  hClose pHandle
  where readJoinRequest hdl = parseText <$> hGetLine hdl
        notifyJoin player@Player{playerId,playerHandle,playerSkilly,playerPos} 
                   bucket@Bucket{bucketPlayers} 
          = do skilly <- atomically $ readTVar playerSkilly
               pos <- atomically $ readTVar playerPos
               let joinEvent = Ejoin playerId skilly pos
               atomically $ notifyBucket playerId bucket joinEvent
               bPlayers <- atomically $ readTVar bucketPlayers
               joinedUsers <- mapM mkJoinedUser $ Set.elems $ Set.delete player bPlayers
               let joinedEvent = Ejoined joinedUsers
               hPutStr playerHandle $ showText joinedEvent
          where mkJoinedUser Player{playerId,playerSkilly,playerPos} = do
                  skilly <- atomically $ readTVar playerSkilly
                  pos <- atomically $ readTVar playerPos
                  return (playerId,skilly,pos)
        stopReadingEvents Session{myEndMark} = do
             -- notifies 'readBucketEvents' process that it has to finish
           atomically $ putTMVar myEndMark ()
             -- check that 'readBucketEvents' has indeed finished
           atomically $ putTMVar myEndMark ()


mkRoomServer :: PortNumber -> Room -> Server
mkRoomServer port room
  = Server {
      serverAddr    = Sock.SockAddrInet port Sock.iNADDR_ANY
    , serverTyp     = Sock.Stream
    , serverRoutine = handlePlayer room
    }


-- * The 'roomd' command

data Roomd = Roomd deriving Typeable

instance Command Roomd (Record Config) where
  run _ conf _ = daemonize $ do
    room <- atomically $ newRoom
    let roomServer = mkRoomServer (fromIntegral $ port conf) room
    roomd <- serveOne Nothing roomServer
    waitFor roomd
    where daemonize | daemon conf = Posix.daemonize
                    | otherwise   = id

main :: IO ()
main = getArgs >>= execute Roomd
